from .builder import DATASETS
from .custom import CustomDataset

'''
background [0, 0, 0]    index:0
build_up [255, 0, 0]    index:1
farmland [0, 255, 0]    index:2
forest [0, 255, 255]    index:3
meadow [255, 255, 0]    index:4
water [0, 0, 255]       index:5
'''
@DATASETS.register_module()
class GIDDataset(CustomDataset):

    CLASSES = (
        'build_up', 'farmland', 'forest', 'meadow', 'water', )

    PALETTE = [[255,0,0],[0,255,0],[0,255,255],[255,255,0],[0,0,255]]

    def __init__(self, **kwargs):
        super(GIDDataset, self).__init__(
            img_suffix='.png',
            seg_map_suffix='.png',
            # add input
            ir_img_suffix='.png',
            dsm_img_suffix='.png',
            reduce_zero_label=True,
            **kwargs)
